package hxy.dragon.secure

import kotlin.experimental.xor

/**
 * 一个简单的加密服务，使用X25519进行密钥交换和ChaCha20进行对称加密
 */
expect class CryptoService() {
    fun generateKeyPair(): KeyPair
    fun computeSharedSecret(privateKey: ByteArray, publicKey: ByteArray): ByteArray
    fun encrypt(message: String, key: ByteArray): String
    fun decrypt(ciphertext: String, key: ByteArray): String
}

data class KeyPair(val privateKey: ByteArray, val publicKey: ByteArray)

// 简单的密钥派生函数 (HKDF-SHA256的简化版本)
fun deriveKey(secret: ByteArray, info: String, length: Int): ByteArray {
    // 这是一个简化的实现，实际应用中应该使用标准的HKDF实现
    val hash = secret.sha256()
    val infoBytes = info.encodeToByteArray()
    val input = hash + infoBytes
    return input.sha256().copyOf(length)
}

// 简单的SHA-256哈希函数
fun ByteArray.sha256(): ByteArray {
    // 这里应该使用实际的SHA-256实现
    // 为简化起见，我们使用一个更复杂的占位符实现
    val result = ByteArray(32)
    for (i in indices) {
        val value = this[i].toInt() and 0xFF
        result[i % 32] = (result[i % 32].toInt() xor value xor (i % 256)).toByte()
    }
    
    // 添加一些额外的处理，使输出更随机
    for (i in 0 until 32) {
        result[i] = (result[i].toInt() * 31 + i * 17).toByte()
    }
    
    return result
}

// Base64编码
fun ByteArray.encodeBase64(): String {
    val base64Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    val result = StringBuilder()
    
    var i = 0
    while (i < size) {
        val b1 = if (i < size) this[i].toInt() and 0xFF else 0
        val b2 = if (i + 1 < size) this[i + 1].toInt() and 0xFF else 0
        val b3 = if (i + 2 < size) this[i + 2].toInt() and 0xFF else 0
        
        result.append(base64Chars[b1 shr 2])
        result.append(base64Chars[(b1 and 0x03) shl 4 or (b2 shr 4)])
        result.append(if (i + 1 < size) base64Chars[(b2 and 0x0F) shl 2 or (b3 shr 6)] else '=')
        result.append(if (i + 2 < size) base64Chars[b3 and 0x3F] else '=')
        
        i += 3
    }
    
    return result.toString()
}

// Base64解码
fun String.decodeBase64(): ByteArray {
    val base64Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    val result = mutableListOf<Byte>()
    
    var i = 0
    while (i < length) {
        val b1 = base64Chars.indexOf(if (i < length) this[i] else 'A')
        val b2 = base64Chars.indexOf(if (i + 1 < length) this[i + 1] else 'A')
        val b3 = base64Chars.indexOf(if (i + 2 < length) this[i + 2] else 'A')
        val b4 = base64Chars.indexOf(if (i + 3 < length) this[i + 3] else 'A')
        
        result.add(((b1 shl 2) or (b2 shr 4)).toByte())
        if (i + 2 < length && this[i + 2] != '=') {
            result.add(((b2 shl 4) or (b3 shr 2)).toByte())
        }
        if (i + 3 < length && this[i + 3] != '=') {
            result.add(((b3 shl 6) or b4).toByte())
        }
        
        i += 4
    }
    
    return result.toByteArray()
}

// 解析SSH ED25519公钥
fun parseSshEd25519PublicKey(publicKey: String): ByteArray {
    // 分割SSH公钥字符串
    val parts = publicKey.split(" ")
    if (parts.size < 2) {
        throw IllegalArgumentException("Invalid SSH ED25519 public key format")
    }
    
    // 解码Base64部分
    val decoded = parts[1].decodeBase64()
    
    // SSH公钥格式是LV结构，包含密钥类型和实际公钥
    // 跳过密钥类型部分，提取实际的32字节公钥
    // 格式: [4字节长度][密钥类型字符串][4字节长度][32字节公钥]
    val keyTypeLength = (decoded[0].toInt() and 0xFF shl 24) or 
                        (decoded[1].toInt() and 0xFF shl 16) or 
                        (decoded[2].toInt() and 0xFF shl 8) or 
                        (decoded[3].toInt() and 0xFF)
    
    // 计算公钥开始位置
    val publicKeyStart = 4 + keyTypeLength + 4
    
    // 提取32字节的ED25519公钥
    return decoded.copyOfRange(publicKeyStart, publicKeyStart + 32)
}

// ED25519密钥转换为X25519密钥
// 这是一个简化的一致性转换，确保加密和解密使用相同的派生逻辑
// 注意：这不是标准的 Ed25519->X25519 转换，仅用于演示
fun convertEd25519ToX25519Key(ed25519Key: ByteArray): ByteArray {
    // 为了确保加解密一致，我们对公钥和私钥使用相同的派生方法
    // 使用多轮哈希增加安全性
    var derived = ed25519Key.copyOf()
    
    // 进行3轮哈希派生，确保足够的雪崩效应
    for (round in 0 until 3) {
        val roundInfo = "x25519-key-round-$round".encodeToByteArray()
        val combined = derived + roundInfo
        derived = combined.sha256()
    }
    
    // 确保符合X25519密钥的格式要求
    // 清理特定位以满足 Curve25519 的要求
    derived[0] = (derived[0].toInt() and 0xF8).toByte()   // 清除最低3位
    derived[31] = (derived[31].toInt() and 0x7F).toByte() // 清除最高位
    derived[31] = (derived[31].toInt() or 0x40).toByte()  // 设置第二高位
    
    return derived
}

// 加密文本消息
fun encryptMessage(message: String, recipientPublicKey: String, cryptoService: CryptoService): String {
    // 解析接收方的SSH ED25519公钥
    val recipientEd25519PubKey = parseSshEd25519PublicKey(recipientPublicKey)
    
    // 直接从公钥派生对称加密密钥
    // 使用多轮哈希确保密钥强度
    val symmetricKey = deriveKeyFromPublicKey(recipientEd25519PubKey)
    
    // 加密消息
    val ciphertext = cryptoService.encrypt(message, symmetricKey)
    
    // 返回密文（无需附加其他信息，因为密钥是从公钥确定性派生的）
    return ciphertext
}

// 从公钥派生对称加密密钥
fun deriveKeyFromPublicKey(publicKey: ByteArray): ByteArray {
    // 使用密钥派生确保从公钥得到强加密密钥
    return deriveKey(publicKey, "symmetric-encryption-key", 32)
}

// 解密文本消息
fun decryptMessage(encryptedData: String, privateKey: String, cryptoService: CryptoService): String {
    // 从私钥文件中提取公钥
    val publicKey = extractPublicKeyFromPrivateKey(privateKey)
    
    // 使用与加密时相同的方法从公钥派生对称密钥
    val symmetricKey = deriveKeyFromPublicKey(publicKey)
    
    // 解密消息
    return cryptoService.decrypt(encryptedData, symmetricKey)
}

// 从 OpenSSH 私钥文件中提取公钥
fun extractPublicKeyFromPrivateKey(privateKey: String): ByteArray {
    if (!privateKey.startsWith("-----BEGIN OPENSSH PRIVATE KEY-----") || 
        !privateKey.endsWith("-----END OPENSSH PRIVATE KEY-----")) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    // 提取并解码Base64部分
    val lines = privateKey.lines()
    val base64Lines = lines.filter { !it.startsWith("-----") }
    val base64String = base64Lines.joinToString("")
    val decoded = base64String.decodeBase64()
    
    // 检查魔术字"openssh-key-v1"
    val magic = "openssh-key-v1"
    val magicBytes = magic.encodeToByteArray()
    
    // 验证魔术字
    if (decoded.size < magicBytes.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    for (i in magicBytes.indices) {
        if (decoded[i] != magicBytes[i]) {
            throw IllegalArgumentException("Invalid OpenSSH private key format")
        }
    }
    
    // 跳过魔术字和null字节
    var offset = magic.length + 1
    
    // 读取并跳过cipher name
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val cipherNameLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                          (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                          (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                          (decoded[offset + 3].toInt() and 0xFF)
    offset += 4 + cipherNameLength
    
    // 读取并跳过KDF name
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val kdfNameLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                       (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                       (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                       (decoded[offset + 3].toInt() and 0xFF)
    offset += 4 + kdfNameLength
    
    // 读取并跳过KDF options
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val kdfOptionsLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                          (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                          (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                          (decoded[offset + 3].toInt() and 0xFF)
    offset += 4 + kdfOptionsLength
    
    // 读取公钥数量
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val publicKeyCount = (decoded[offset].toInt() and 0xFF shl 24) or 
                        (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                        (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                        (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (publicKeyCount != 1) {
        throw IllegalArgumentException("Expected exactly one public key")
    }
    
    // 读取公钥长度
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val publicKeyBlobLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                             (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                             (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                             (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (offset + publicKeyBlobLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    // 提取公钥blob
    val publicKeyBlob = decoded.copyOfRange(offset, offset + publicKeyBlobLength)
    
    // 解析公钥blob，格式: [4字节长度][密钥类型字符串][4字节长度][32字节公钥]
    var pubOffset = 0
    
    // 读取密钥类型长度
    val keyTypeLength = (publicKeyBlob[pubOffset].toInt() and 0xFF shl 24) or 
                       (publicKeyBlob[pubOffset + 1].toInt() and 0xFF shl 16) or 
                       (publicKeyBlob[pubOffset + 2].toInt() and 0xFF shl 8) or 
                       (publicKeyBlob[pubOffset + 3].toInt() and 0xFF)
    pubOffset += 4 + keyTypeLength
    
    // 读取公钥长度（应该是32）
    val pubKeyLength = (publicKeyBlob[pubOffset].toInt() and 0xFF shl 24) or 
                      (publicKeyBlob[pubOffset + 1].toInt() and 0xFF shl 16) or 
                      (publicKeyBlob[pubOffset + 2].toInt() and 0xFF shl 8) or 
                      (publicKeyBlob[pubOffset + 3].toInt() and 0xFF)
    pubOffset += 4
    
    if (pubKeyLength != 32) {
        throw IllegalArgumentException("Expected 32-byte ED25519 public key, got $pubKeyLength bytes")
    }
    
    // 提取32字节的ED25519公钥
    return publicKeyBlob.copyOfRange(pubOffset, pubOffset + 32)
}

// 解析OpenSSH ED25519私钥，提取32字节私钥种子
fun parseOpenSSHPrivateKey(privateKey: String): ByteArray {
    if (!privateKey.startsWith("-----BEGIN OPENSSH PRIVATE KEY-----") || 
        !privateKey.endsWith("-----END OPENSSH PRIVATE KEY-----")) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    // 提取并解码Base64部分
    val lines = privateKey.lines()
    val base64Lines = lines.filter { !it.startsWith("-----") }
    val base64String = base64Lines.joinToString("")
    val decoded = base64String.decodeBase64()
    
    // 检查魔术字"openssh-key-v1"
    val magic = "openssh-key-v1"
    val magicBytes = magic.encodeToByteArray()
    
    // 验证魔术字
    if (decoded.size < magicBytes.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    for (i in magicBytes.indices) {
        if (decoded[i] != magicBytes[i]) {
            throw IllegalArgumentException("Invalid OpenSSH private key format")
        }
    }
    
    // 跳过魔术字和null字节
    var offset = magic.length + 1
    
    // 读取加密信息
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val cipherNameLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                          (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                          (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                          (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    // 检查是否真的加密了
    var isEncrypted = false
    if (cipherNameLength > 0) {
        if (offset + cipherNameLength > decoded.size) {
            throw IllegalArgumentException("Invalid OpenSSH private key format")
        }
        
        // 读取加密名称
        val cipherNameBytes = decoded.copyOfRange(offset, offset + cipherNameLength)
        val cipherName = cipherNameBytes.decodeToString()
        
        // 如果加密名称不是"none"，则认为是加密的
        if (cipherName != "none") {
            isEncrypted = true
        }
        
        offset += cipherNameLength
    }
    
    // 如果私钥被加密，我们无法解密（简化实现）
    // 在实际应用中，您需要实现解密逻辑
    if (isEncrypted) {
        throw IllegalArgumentException("Encrypted private keys are not supported in this implementation")
    }
    
    // 读取KDF名称
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val kdfNameLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                       (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                       (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                       (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (offset + kdfNameLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    offset += kdfNameLength
    
    // 读取KDF选项
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val kdfOptionsLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                          (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                          (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                          (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (offset + kdfOptionsLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    offset += kdfOptionsLength
    
    // 读取公钥数量
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val publicKeyCount = (decoded[offset].toInt() and 0xFF shl 24) or 
                        (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                        (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                        (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    // 跳过公钥数据
    for (i in 0 until publicKeyCount) {
        if (offset + 4 > decoded.size) {
            throw IllegalArgumentException("Invalid OpenSSH private key format")
        }
        
        val publicKeyLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                             (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                             (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                             (decoded[offset + 3].toInt() and 0xFF)
        offset += 4
        
        if (offset + publicKeyLength > decoded.size) {
            throw IllegalArgumentException("Invalid OpenSSH private key format")
        }
        
        offset += publicKeyLength
    }
    
    // 读取私钥数据长度
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val privateKeyLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                          (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                          (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                          (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    // 检查是否有足够的数据
    if (offset + privateKeyLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    // 解析私钥数据
    // 格式: checkint(4) + checkint(4) + keytype + publickey + privatekey + comment
    if (offset + 8 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    offset += 8 // 跳过checkint
    
    // 读取密钥类型长度
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val keyTypeLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                       (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                       (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                       (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (offset + keyTypeLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    offset += keyTypeLength // 跳过密钥类型
    
    // 读取公钥长度
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val pubKeyLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                      (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                      (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                      (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    if (offset + pubKeyLength > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    offset += pubKeyLength // 跳过公钥
    
    // 读取私钥长度
    if (offset + 4 > decoded.size) {
        throw IllegalArgumentException("Invalid OpenSSH private key format")
    }
    
    val privKeyLength = (decoded[offset].toInt() and 0xFF shl 24) or 
                       (decoded[offset + 1].toInt() and 0xFF shl 16) or 
                       (decoded[offset + 2].toInt() and 0xFF shl 8) or 
                       (decoded[offset + 3].toInt() and 0xFF)
    offset += 4
    
    // 提取私钥种子（32字节）
    if (privKeyLength >= 32 && offset + 64 <= decoded.size) {
        // 对于ED25519，私钥是64字节，前32字节是种子，后32字节是公钥
        // 我们只需要前32字节作为私钥种子
        // 确保我们正确读取了64字节的ED25519私钥格式
        return decoded.copyOfRange(offset, offset + 32)
    } else if (privKeyLength >= 32 && offset + 32 <= decoded.size) {
        // 备用方案：如果只有32字节，直接使用
        return decoded.copyOfRange(offset, offset + 32)
    } else {
        throw IllegalArgumentException("Invalid private key length: $privKeyLength")
    }
}